#!/usr/bin/env python
import os
import sys
import subprocess
import base64
import argparse
import mimetypes
# import json
# import logging
from PIL import Image
from concurrent.futures import ThreadPoolExecutor, as_completed

# # 配置日志系统
# logging.basicConfig(
#     level=logging.INFO,
#     format='%(asctime)s - %(levelname)s - %(message)s',
#     handlers=[
#         logging.StreamHandler(sys.stdout),
#         logging.FileHandler('compression.log')
#     ]
# )
# logger = logging.getLogger(__name__)

# 默认压缩配置
DEFAULT_COMPRESSION_CONFIGS = {
    "ffmpeg": {
        "default": {
            "crf": 28,
            "preset": "slow",
            "audio_bitrate": "64k",
            "scale": "40:40",
            "quality": 80
        },
        "high_quality_video": {
            "crf": 23,
            "preset": "medium",
            "audio_bitrate": "128k"
        },
        "low_quality": {
            "crf": 35,
            "preset": "ultrafast",
            "audio_bitrate": "32k",
            "scale": "20:20",
            "quality": 50
        }
    },
    "pillow": {
        "default": {
            "quality": 80,
            "optimize": True
        },
        "high_quality": {
            "quality": 95,
            "optimize": True
        },
        "low_quality": {
            "quality": 50,
            "optimize": True
        }
    },
    "pngquant": {
        "default": {
            "quality": "65-80",
            "speed": 3
        },
        "high_quality": {
            "quality": "90-100",
            "speed": 1
        },
        "low_quality": {
            "quality": "40-60",
            "speed": 5
        }
    },
    "jpegoptim": {
        "default": {
            "max_quality": 80
        },
        "high_quality": {
            "max_quality": 95
        },
        "low_quality": {
            "max_quality": 50
        }
    },
    "webp": {
        "default": {
            "quality": 80,
            "method": 4
        },
        "high_quality": {
            "quality": 95,
            "method": 4
        },
        "low_quality": {
            "quality": 50,
            "method": 1
        }
    }
}

# 文件类型到压缩工具的映射
FILE_TOOL_MAPPING = {
    "video": ["ffmpeg"],
    "image/jpeg": ["jpegoptim", "ffmpeg", "pillow", "webp"],
    "image/png": ["pngquant", "ffmpeg", "pillow", "webp"],
    "image/webp": ["ffmpeg", "pillow", "webp"],
    "image/gif": ["ffmpeg", "pillow"],
    "image/bmp": ["ffmpeg", "pillow"],
    "image/tiff": ["ffmpeg", "pillow"]
}

def save_base64_encoding(input_file, output_dir, filename_no_ext, extension):
    """保存文件的base64编码"""
    base64_output_file = os.path.join(output_dir, f"{filename_no_ext}.txt")
    try:
        with open(input_file, 'rb') as f:
            encoded_string = base64.b64encode(f.read()).decode('utf-8')
        
        with open(base64_output_file, 'w') as f:
            f.write(f"data:{mimetypes.guess_type(input_file)[0]};base64,{encoded_string}")
        # logger.info(f"Base64编码已保存到: {base64_output_file}")
        return True
    except Exception as e:
        # logger.error(f"生成Base64编码失败 {input_file}: {e}")
        return False

def compress_video_ffmpeg(input_file, output_dir, filename_no_ext, extension, config):
    """使用ffmpeg压缩视频"""
    output_file = os.path.join(output_dir, f"{filename_no_ext}.{extension}")
    # logger.info(f"正在压缩视频: {input_file}")
    
    # 构建ffmpeg命令
    command = [
        'ffmpeg', '-y', '-i', input_file,
        '-c:v', 'libx264', '-crf', str(config.get('crf', 28)),
        '-preset', config.get('preset', 'slow'),
        '-b:a', config.get('audio_bitrate', '64k'),
        '-movflags', '+faststart',
        '-vf', f"scale={config.get('scale', 'trunc(iw/2)*2:trunc(ih/2)*2')}",
        output_file
    ]
    
    try:
        result = subprocess.run(
            command, 
            check=True, 
            stdout=subprocess.PIPE, 
            stderr=subprocess.STDOUT,
            text=True
        )
        # logger.info(f"视频压缩成功: {os.path.basename(input_file)}")
        # logger.debug(f"FFmpeg输出:\n{result.stdout}")
        
        # 保存base64编码
        save_base64_encoding(output_file, output_dir, filename_no_ext, extension)
        return True
    except subprocess.CalledProcessError as e:
        # logger.error(f"视频压缩失败 {input_file}: {e}\n输出: {e.stdout}")
        return False

def compress_image_ffmpeg(input_file, output_dir, filename_no_ext, extension, config):
    """使用ffmpeg压缩图片"""
    output_file = os.path.join(output_dir, f"{filename_no_ext}.{extension}")
    # logger.info(f"正在压缩图片: {input_file}")
    
    command = [
        'ffmpeg', '-y', '-i', input_file,
        '-vf', f"scale={config.get('scale', '40:40')}",
        '-q:v', str(config.get('quality', 80)),
        output_file
    ]
    
    try:
        result = subprocess.run(
            command, 
            check=True, 
            stdout=subprocess.PIPE, 
            stderr=subprocess.STDOUT,
            text=True
        )
        # logger.info(f"图片压缩成功: {os.path.basename(input_file)}")
        # logger.debug(f"FFmpeg输出:\n{result.stdout}")
        
        # 保存base64编码
        save_base64_encoding(output_file, output_dir, filename_no_ext, extension)
        return True
    except subprocess.CalledProcessError as e:
        # logger.error(f"图片压缩失败 {input_file}: {e}\n输出: {e.stdout}")
        return False

def compress_image_pillow(input_file, output_dir, filename_no_ext, extension, config):
    """使用Pillow压缩图片"""
    output_file = os.path.join(output_dir, f"{filename_no_ext}.{extension}")
    # logger.info(f"正在使用Pillow压缩图片: {input_file}")
    
    try:
        with Image.open(input_file) as img:
            # 自动旋转图片
            # if hasattr(img, '_getexif'):
            #     exif = img._getexif()
            #     if exif:
            #         orientation = exif.get(0x0112)
            #         if orientation == 3:
            #             img = img.rotate(180, expand=True)
            #         elif orientation == 6:
            #             img = img.rotate(270, expand=True)
            #         elif orientation == 8:
            #             img = img.rotate(90, expand=True)
            
            # 调整大小（如果需要）
            if 'scale' in config:
                width, height = map(int, config['scale'].split(':'))
                img = img.resize((width, height), Image.Resampling.LANCZOS)
            
            # 保存图片
            save_args = {
                'optimize': config.get('optimize', True),
                'quality': config.get('quality', 80)
            }
            
            if extension.lower() == 'webp':
                save_args['method'] = config.get('method', 4)
            
            img.save(output_file, **save_args)
        
        # logger.info(f"Pillow图片压缩成功: {os.path.basename(input_file)}")
        
        # 保存base64编码
        save_base64_encoding(output_file, output_dir, filename_no_ext, extension)
        return True
    except Exception as e:
        # logger.error(f"Pillow图片压缩失败 {input_file}: {e}")
        return False

def compress_image_pngquant(input_file, output_dir, filename_no_ext, extension, config):
    """使用pngquant压缩PNG图片"""
    output_file = os.path.join(output_dir, f"{filename_no_ext}.{extension}")
    # logger.info(f"正在使用pngquant压缩图片: {input_file}")
    
    command = [
        'pngquant',
        '--quality', config.get('quality', '65-80'),
        '--speed', str(config.get('speed', 3)),
        '--output', output_file,
        '--force',
        '--skip-if-larger',
        input_file
    ]
    
    try:
        result = subprocess.run(
            command, 
            check=True, 
            stdout=subprocess.PIPE, 
            stderr=subprocess.STDOUT,
            text=True
        )
        # logger.info(f"pngquant图片压缩成功: {os.path.basename(input_file)}")
        # logger.debug(f"pngquant输出:\n{result.stdout}")
        
        # 保存base64编码
        save_base64_encoding(output_file, output_dir, filename_no_ext, extension)
        return True
    except FileNotFoundError:
        print("错误: pngquant未安装或不在PATH中。请访问 https://pngquant.org/ 安装。")
        # logger.error("错误: pngquant未安装或不在PATH中。请访问 https://pngquant.org/ 安装。")
        return False
    except subprocess.CalledProcessError as e:
        print(f"pngquant图片压缩失败 {input_file}: {e}\n输出: {e.stdout}")
        # logger.error(f"pngquant图片压缩失败 {input_file}: {e}\n输出: {e.stdout}")
        return False

def compress_image_jpegoptim(input_file, output_dir, filename_no_ext, extension, config):
    """使用jpegoptim压缩JPEG图片"""
    output_file = os.path.join(output_dir, f"{filename_no_ext}.{extension}")
    # logger.info(f"正在使用jpegoptim压缩图片: {input_file}")
    
    # 复制文件到输出目录
    try:
        import shutil
        shutil.copy(input_file, output_file)
    except Exception as e:
        # logger.error(f"文件复制失败 {input_file}: {e}")
        return False

    command = [
        'jpegoptim',
        '--max', str(config.get('max_quality', 80)),
        '--strip-all',
        '--overwrite',
        output_file
    ]
    
    try:
        result = subprocess.run(
            command, 
            check=True, 
            stdout=subprocess.PIPE, 
            stderr=subprocess.STDOUT,
            text=True
        )
        # logger.info(f"jpegoptim图片压缩成功: {os.path.basename(input_file)}")
        # logger.debug(f"jpegoptim输出:\n{result.stdout}")
        
        # 保存base64编码
        save_base64_encoding(output_file, output_dir, filename_no_ext, extension)
        return True
    except FileNotFoundError:
        # logger.error("错误: jpegoptim未安装或不在PATH中。请访问 https://github.com/tjko/jpegoptim 安装。")
        return False
    except subprocess.CalledProcessError as e:
        # logger.error(f"jpegoptim图片压缩失败 {input_file}: {e}\n输出: {e.stdout}")
        return False

def compress_webp(input_file, output_dir, filename_no_ext, config):
    """使用Pillow压缩为WebP格式"""
    output_file = os.path.join(output_dir, f"{filename_no_ext}.webp")
    # logger.info(f"正在转换为WebP格式: {input_file}")
    
    try:
        with Image.open(input_file) as img:
            # 自动旋转图片
            # if hasattr(img, '_getexif'):
            #     exif = img._getexif()
            #     if exif:
            #         orientation = exif.get(0x0112)
            #         if orientation == 3:
            #             img = img.rotate(180, expand=True)
            #         elif orientation == 6:
            #             img = img.rotate(270, expand=True)
            #         elif orientation == 8:
            #             img = img.rotate(90, expand=True)
            
            # 调整大小（如果需要）
            if 'scale' in config:
                width, height = map(int, config['scale'].split(':'))
                img = img.resize((width, height), Image.Resampling.LANCZOS)
            
            # 保存为WebP
            img.save(
                output_file, 
                'WEBP', 
                quality=config.get('quality', 80),
                method=config.get('method', 4)
            )
        
        # logger.info(f"WebP转换成功: {os.path.basename(input_file)}")
        
        # 保存base64编码
        save_base64_encoding(output_file, output_dir, filename_no_ext, "webp")
        return True
    except Exception as e:
        # logger.error(f"WebP转换失败 {input_file}: {e}")
        return False

# def load_config(config_path='config.json'):
#     """从配置文件加载压缩配置"""
#     if os.path.exists(config_path):
#         try:
#             with open(config_path, 'r') as f:
#                 user_config = json.load(f)
            
#             # 合并默认配置和用户配置
#             config = DEFAULT_COMPRESSION_CONFIGS.copy()
#             for tool, profiles in user_config.get('compression_configs', {}).items():
#                 if tool in config:
#                     config[tool].update(profiles)
#                 else:
#                     config[tool] = profiles
            
#             # 更新文件类型到工具的映射
#             file_mapping = user_config.get('file_tool_mapping', {})
#             for mime_type, tools in file_mapping.items():
#                 FILE_TOOL_MAPPING[mime_type] = tools
            
#             return config
#         except (json.JSONDecodeError, KeyError) as e:
#             # logger.error(f"配置文件错误: {e}，使用默认配置")
#     return DEFAULT_COMPRESSION_CONFIGS

def compress_file(file_path, output_dir, config, compress_tool=None, tool_config_key='default', threads=1):
    """压缩单个文件"""
    filename = os.path.basename(file_path)
    name, ext = os.path.splitext(filename)
    extension = ext.lstrip('.').lower()
    
    # 获取文件MIME类型
    mime_type, _ = mimetypes.guess_type(file_path)
    if mime_type is None:
        # logger.warning(f"无法确定文件类型: {filename}，跳过")
        return False
    
    # 确定文件类别（视频或图片）
    file_category = "video" if mime_type.startswith("video/") else mime_type
    
    # 获取合适的压缩工具
    supported_tools = FILE_TOOL_MAPPING.get(file_category, [])
    # print(supported_tools)
    if not supported_tools:
        # logger.warning(f"不支持的文件类型: {mime_type}，跳过 {filename}")
        return False
    
    # 如果未指定工具，使用第一个支持的压缩工具
    if compress_tool is None or compress_tool not in supported_tools:
        compress_tool = supported_tools[0]
        # logger.info(f"自动选择工具 '{compress_tool}' 处理 {filename}")
    # print(compress_tool)
    # 获取工具配置
    tool_config = config.get(compress_tool, {}).get(tool_config_key, {})
    if not tool_config:
        # logger.warning(f"未找到配置 '{tool_config_key}' 对于工具 '{compress_tool}'，使用默认配置")
        tool_config = config.get(compress_tool, {}).get('default', {})
    # print((file_path, output_dir, name, extension, tool_config))
    # 执行压缩
    if compress_tool == 'ffmpeg':
        if file_category == "video":
            return compress_video_ffmpeg(file_path, output_dir, name, extension, tool_config)
        else:
            return compress_image_ffmpeg(file_path, output_dir, name, extension, tool_config)
    elif compress_tool == 'pillow':
        return compress_image_pillow(file_path, output_dir, name, extension, tool_config)
    elif compress_tool == 'pngquant':
        return compress_image_pngquant(file_path, output_dir, name, extension, tool_config)
    elif compress_tool == 'jpegoptim':
        return compress_image_jpegoptim(file_path, output_dir, name, extension, tool_config)
    elif compress_tool == 'webp':
        return compress_webp(file_path, output_dir, name, tool_config)
    else:
        # logger.error(f"不支持的压缩工具: {compress_tool}")
        return False

def main():
    parser = argparse.ArgumentParser(description='压缩视频和图片文件。')
    parser.add_argument('input_directory', type=str, help='包含要压缩文件的输入目录。')
    parser.add_argument('--tool', type=str, default=None, help='选择压缩工具 (ffmpeg, pillow, pngquant, jpegoptim, webp)。')
    parser.add_argument('--config', type=str, default='config.json', help='配置文件路径。')
    parser.add_argument('--config_key', type=str, default='default', help='选择配置中工具的配置键。')
    parser.add_argument('--threads', type=int, default=4, help='并行处理线程数。')
    parser.add_argument('--debug', action='store_true', help='启用调试模式。')
    args = parser.parse_args()
    
    # if args.debug:
        # logger.setLevel(logging.DEBUG)
        # logger.debug("调试模式已启用")
    
    # 加载配置
    # config = load_config(args.config)
    
    input_dir = args.input_directory
    output_dir = os.path.join(input_dir, "py_output")
    
    if not os.path.isdir(input_dir):
        # logger.error(f"输入目录 '{input_dir}' 不存在或不是一个目录。")
        sys.exit(1)
    
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        # logger.info(f"创建输出目录: {output_dir}")
    
    # 收集所有文件
    files_to_compress = []
    for filename in os.listdir(input_dir):
        file_path = os.path.join(input_dir, filename)
        if os.path.isfile(file_path):
            files_to_compress.append(file_path)
    
    if not files_to_compress:
        # logger.warning("未找到可压缩的文件")
        return
    
    # logger.info(f"找到 {len(files_to_compress)} 个文件需要压缩")
    
    # 并行处理文件
    success_count = 0
    with ThreadPoolExecutor(max_workers=args.threads) as executor:
        futures = {}
        for file_path in files_to_compress:
            future = executor.submit(
                compress_file,
                file_path,
                output_dir,
                DEFAULT_COMPRESSION_CONFIGS,
                args.tool,
                args.config_key,
                args.threads
            )
            futures[future] = file_path
        
        for future in as_completed(futures):
            file_path = futures[future]
            try:
                if future.result():
                    success_count += 1
                    # logger.info(f"成功压缩: {os.path.basename(file_path)}")
                # else:
                    # logger.warning(f"压缩失败: {os.path.basename(file_path)}")
            except Exception as e:
                print(f"处理文件时发生错误 {file_path}: {e}")
                # logger.error(f"处理文件时发生错误 {file_path}: {e}")
    
    # logger.info(f"压缩完成! 成功: {success_count}/{len(files_to_compress)}")

if __name__ == "__main__":
    mimetypes.init()
    main()